# -*- coding:utf-8 -*-

import numpy as np
import cv2 as cv
from awmf import Awmf
from damf import Damf
from nafsmf import Nafsmf
from ehga import Ehga
from image_evaluate import ImageEvalue


def add_SAP(img_in, SNR):
    '''
    为灰度图像添加椒盐噪声
    :param SNR:信噪比
    '''
    img_out = img_in.copy()
    h, w = img_out.shape
    mask = np.random.choice((0, 1, 2), size=(h, w), p=[SNR, (1 - SNR) / 2., (1 - SNR) / 2.])
    img_out[mask == 1] = 255  # 盐噪声
    img_out[mask == 2] = 0  # 椒噪声
    return img_out


def prepare_noise_image(img_path, SAP):
    # 读取图像
    im = cv.imread(img_path, 0)
    im = np.array(im)
    # 添加噪声
    noise_image = add_SAP(im, SAP)
    # 保存噪声图
    cv.imwrite('./noise{}.tiff'.format(int(100 - SAP * 100)), noise_image)
    return noise_image


def process_and_save_image(noise_image_path, SAP):
    # 滤波器初始化
    noise_image = cv.imread(noise_image_path, 0)
    awm_filter = Awmf(h=1, w_max=39)
    dam_filter = Damf()
    nafsm_filter = Nafsmf(h=1, s_max=3, T1=10, T2=30)
    ehga_filter = Ehga(pop_num=60, crossover_num=57, mutation_rate=0.02, lambda_=0.075, beta=1, iter_max=5)
    # 处理图像并保存图像
    cv.imshow('noise image', noise_image)
    cv.waitKey(0)
    print('Damf')
    damf_denoise_image = dam_filter.process_image(noise_image)
    cv.imwrite('./damf_denoise_image{}.tiff'.format(int(100 - SAP * 100)), damf_denoise_image)
    cv.imshow('Damf', damf_denoise_image)
    cv.waitKey(0)

    print('Nafsmf')
    nafsmf_denoise_image = nafsm_filter.process_image(noise_image)
    cv.imwrite('./nafsmf_denoise_image{}.tiff'.format(int(100 - SAP * 100)), nafsmf_denoise_image)
    cv.imshow('Nafsmf', nafsmf_denoise_image)
    cv.waitKey(0)

    print('Awmf')
    awmf_denoise_image = awm_filter.process_image(noise_image)
    cv.imwrite('./awmf_denoise_image{}.tiff'.format(int(100 - SAP * 100)), awmf_denoise_image)
    cv.imshow('Awmf', awmf_denoise_image)
    cv.waitKey(0)

    print('Ehga')
    ehga_denoise_image = ehga_filter.process_image(noise_image)
    cv.imwrite('./ehga_denoise_image{}.tiff'.format(int(100 - SAP * 100)), ehga_denoise_image)
    cv.imshow('Ehga', ehga_denoise_image)
    cv.waitKey(0)


def evaluate_image(origin_img_path, denoise_image_path, noise_image_path):
    origin_img = cv.imread(origin_img_path, 0)
    denoise_image = cv.imread(denoise_image_path, 0)
    noise_image = cv.imread(noise_image_path, 0)
    evaluator = ImageEvalue()
    psnr = evaluator.PSNR(origin_img, denoise_image)
    ssim = evaluator.SSIM(origin_img, denoise_image)
    ief = evaluator.IEF(origin_img, denoise_image, noise_image)
    uqi = evaluator.UQI(origin_img, denoise_image)
    print('psnr:', psnr)
    print('ssim', ssim)
    print('ief', ief)
    print('uqi', uqi)


if __name__ == '__main__':
    origin_img_path = 'lena.tiff'
    noise_image_path = 'noise90.tiff'
    denoise_image_path = 'ehga_denoise_image90.tiff'
    SAP = 0.1
    noise_image = prepare_noise_image(origin_img_path, SAP)
    process_and_save_image(noise_image_path, SAP)
    evaluate_image(origin_img_path, denoise_image_path, noise_image_path)
